// pkg/metadata/store.go
package metadata

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/cockroachdb/pebble"
	// 可选: "github.com/tecbot/gorocksdb"
	// 可选: "go.etcd.io/bbolt"
)

// ChunkInfo 块信息
type ChunkInfo struct {
	Digest   string `json:"digest"`   // sha256 hash
	Size     int64  `json:"size"`     // 大小
	Offset   int64  `json:"offset"`   // 在数据块中的偏移
	RefCount int32  `json:"refcount"` // 引用计数
}

// FileMapping 文件到块的映射
type FileMapping struct {
	Path   string      `json:"path"`
	Mode   int32       `json:"mode"`
	Size   int64       `json:"size"`
	Chunks []ChunkInfo `json:"chunks"`
}

// SnapshotMetadata 快照元数据
type SnapshotMetadata struct {
	Key       string        `json:"key"`
	Parent    string        `json:"parent"`
	Files     []FileMapping `json:"files"`
	CreatedAt int64         `json:"created_at"`
	Size      int64         `json:"size"`
}

// Store 元数据存储接口
type Store interface {
	// 保存快照元数据
	SaveSnapshot(ctx context.Context, meta *SnapshotMetadata) error
	// 加载快照元数据
	LoadSnapshot(ctx context.Context, key string) (*SnapshotMetadata, error)
	// 删除快照元数据
	DeleteSnapshot(ctx context.Context, key string) error

	// 保存块信息
	SaveChunk(ctx context.Context, digest string, info *ChunkInfo) error
	// 加载块信息
	LoadChunk(ctx context.Context, digest string) (*ChunkInfo, error)
	// 列出所有块
	ListChunks(ctx context.Context) ([]*ChunkInfo, error)

	// 保存文件映射
	SaveFileMapping(ctx context.Context, snapshotKey string, mapping *FileMapping) error
	// 加载文件映射
	LoadFileMapping(ctx context.Context, snapshotKey, path string) (*FileMapping, error)

	// 递增块引用计数
	IncrementChunkRef(ctx context.Context, digest string) error
	// 递减块引用计数
	DecrementChunkRef(ctx context.Context, digest string) error

	// 关闭存储
	Close() error
}

// PebbleStore 基于 Pebble 的实现
type PebbleStore struct {
	db *pebble.DB
	mu sync.RWMutex
}

// NewStore 创建元数据存储
func NewStore(backend, rootDir string) (Store, error) {
	switch backend {
	case "pebble", "rocksdb":
		return NewPebbleStore(rootDir)
	// 可添加 BoltDB 实现
	default:
		return nil, fmt.Errorf("unsupported backend: %s", backend)
	}
}

// NewPebbleStore 创建 Pebble 存储
func NewPebbleStore(rootDir string) (*PebbleStore, error) {
	db, err := pebble.Open(rootDir, &pebble.Options{})
	if err != nil {
		return nil, fmt.Errorf("failed to open pebble db: %w", err)
	}
	return &PebbleStore{db: db}, nil
}

// SaveSnapshot 保存快照
func (s *PebbleStore) SaveSnapshot(ctx context.Context, meta *SnapshotMetadata) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	data, err := json.Marshal(meta)
	if err != nil {
		return err
	}

	key := fmt.Sprintf("snapshot:%s", meta.Key)
	return s.db.Set([]byte(key), data, &pebble.WriteOptions{})
}

// LoadSnapshot 加载快照
func (s *PebbleStore) LoadSnapshot(ctx context.Context, key string) (*SnapshotMetadata, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	dbKey := fmt.Sprintf("snapshot:%s", key)
	data, closer, err := s.db.Get([]byte(dbKey))
	if err != nil {
		return nil, err
	}
	defer closer.Close()

	var meta SnapshotMetadata
	if err := json.Unmarshal(data, &meta); err != nil {
		return nil, err
	}
	return &meta, nil
}

// DeleteSnapshot 删除快照
func (s *PebbleStore) DeleteSnapshot(ctx context.Context, key string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	dbKey := fmt.Sprintf("snapshot:%s", key)
	return s.db.Delete([]byte(dbKey), &pebble.WriteOptions{})
}

// SaveChunk 保存块信息
func (s *PebbleStore) SaveChunk(ctx context.Context, digest string, info *ChunkInfo) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	data, err := json.Marshal(info)
	if err != nil {
		return err
	}

	key := fmt.Sprintf("chunk:%s", digest)
	return s.db.Set([]byte(key), data, &pebble.WriteOptions{})
}

// LoadChunk 加载块信息
func (s *PebbleStore) LoadChunk(ctx context.Context, digest string) (*ChunkInfo, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	key := fmt.Sprintf("chunk:%s", digest)
	data, closer, err := s.db.Get([]byte(key))
	if err != nil {
		return nil, err
	}
	defer closer.Close()

	var info ChunkInfo
	if err := json.Unmarshal(data, &info); err != nil {
		return nil, err
	}
	return &info, nil
}

// ListChunks 列出所有块
func (s *PebbleStore) ListChunks(ctx context.Context) ([]*ChunkInfo, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var chunks []*ChunkInfo
	iter, err := s.db.NewIter(&pebble.IterOptions{})
	if err != nil {
		return nil, err
	}
	defer iter.Close()

	for iter.First(); iter.Valid(); iter.Next() {
		key := string(iter.Key())
		if len(key) > 6 && key[:6] == "chunk:" {
			var info ChunkInfo
			if err := json.Unmarshal(iter.Value(), &info); err != nil {
				continue
			}
			chunks = append(chunks, &info)
		}
	}
	return chunks, nil
}

// SaveFileMapping 保存文件映射
func (s *PebbleStore) SaveFileMapping(ctx context.Context, snapshotKey string, mapping *FileMapping) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	data, err := json.Marshal(mapping)
	if err != nil {
		return err
	}

	key := fmt.Sprintf("file:%s:%s", snapshotKey, mapping.Path)
	return s.db.Set([]byte(key), data, &pebble.WriteOptions{})
}

// LoadFileMapping 加载文件映射
func (s *PebbleStore) LoadFileMapping(ctx context.Context, snapshotKey, path string) (*FileMapping, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	key := fmt.Sprintf("file:%s:%s", snapshotKey, path)
	data, closer, err := s.db.Get([]byte(key))
	if err != nil {
		return nil, err
	}
	defer closer.Close()

	var mapping FileMapping
	if err := json.Unmarshal(data, &mapping); err != nil {
		return nil, err
	}
	return &mapping, nil
}

// IncrementChunkRef 递增块引用计数
func (s *PebbleStore) IncrementChunkRef(ctx context.Context, digest string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	key := fmt.Sprintf("chunk:%s", digest)
	data, closer, err := s.db.Get([]byte(key))
	if err != nil {
		return err
	}
	defer closer.Close()

	var info ChunkInfo
	if err := json.Unmarshal(data, &info); err != nil {
		return err
	}

	info.RefCount++
	newData, _ := json.Marshal(info)
	return s.db.Set([]byte(key), newData, &pebble.WriteOptions{})
}

// DecrementChunkRef 递减块引用计数
func (s *PebbleStore) DecrementChunkRef(ctx context.Context, digest string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	key := fmt.Sprintf("chunk:%s", digest)
	data, closer, err := s.db.Get([]byte(key))
	if err != nil {
		return err
	}
	defer closer.Close()

	var info ChunkInfo
	if err := json.Unmarshal(data, &info); err != nil {
		return err
	}

	if info.RefCount > 0 {
		info.RefCount--
	}
	newData, _ := json.Marshal(info)
	return s.db.Set([]byte(key), newData, &pebble.WriteOptions{})
}

// Close 关闭存储
func (s *PebbleStore) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.db.Close()
}
